"""6.6 【上下文类】Context 大模型回调客户端开发"""

import asyncio
import json
from contextlib import AsyncExitStack
from typing import Optional

from mcp import ClientSession
from mcp.server.fastmcp import Context
from mcp.types import CreateMessageRequestParams, CreateMessageResult, TextContent
from openai import OpenAI
from mcp.client.sse import sse_client


class MCPClient:

    def __init__(self):
        #创建线程管理栈
        self.exit_stack = AsyncExitStack()
        self.session: Optional[ClientSession] = None
        self.deepseek = OpenAI(
            api_key="sk-15af4e21f828460683b16ce9e78b2346",
            base_url="https://api.deepseek.com"
        )
    async def sampling_handler(self,context:Context,params:CreateMessageRequestParams)->CreateMessageResult:
        print(f"context: {context}")
        print(f"params: {params}")
        messages=[
            {
                "role":message.role,
                "content":message.content.text
            } for message in params.messages
        ]
        # 调用大模型
        deepseek_response = self.deepseek.chat.completions.create(
            model="deepseek-chat",
            messages=messages,
        )
        # 获取大模型的决策结果
        print("==== deepseek 决策结果：", deepseek_response)
        choice_result = deepseek_response.choices[0]
        return CreateMessageResult(
            role="assistant",
            content=TextContent(text=choice_result.message.content, type="text"),
            model="deepseek-chat"
        )
    """创建连接服务端"""
    async def connect_to_server(self,server_path: str):

        #  创建sse_client
        client = sse_client(url=server_path)
        stdio_transport = await self.exit_stack.enter_async_context(client)
        read_stream, write_stream = stdio_transport

        # 三、创建ClientSession
        client_session = ClientSession(
            read_stream,
            write_stream,
            sampling_callback=self.sampling_handler,

        )
        self.session = await self.exit_stack.enter_async_context(client_session)
        # 四、初始化session
        await self.session.initialize()
        #五、列出所有可用的工具
        response = await self.session.list_tools()
        tools = response.tools
        tool = tools[0]
        print("\nConnected to server with tools:",tools)
        result = await self.session.call_tool(tool.name)
        print("-----------调用工具返回结果：",result)




    #关闭连接
    async def cleanup(self):
        await self.exit_stack.aclose()

async def main():
    client = MCPClient()
    try:
        await client.connect_to_server("http://localhost:8000/sse")

    except Exception as e:
        print(f"Error: {str(e)}")
    finally:
        await client.cleanup()

if __name__ == "__main__":
    asyncio.run(main())